#include "rescue.hpp"
#include "../utils/utils.hpp"
#include "shuf.hpp"

unsigned int MaxPeople = 10;

static const double Eps = 1e-5;

Rescue::Rescue(const Map &map, const std::vector<int> &p, int h, double d) :
	southonly(false), westonly(false),
	m(map), visited(m.nodes.size(), false), people(p), reported(p.size(), false),
	loc(h), time(0), home(h), deadline(d) {

	visited.at(loc) = true;
}

Rescue::Rescue(const Map& map, const std::vector<int>& p, int h, int c, double d)  :
	southonly(false), westonly(false),
	m(map), visited(m.nodes.size(), false), people(p), reported(p.size(), false),
	loc(c), time(0), home(h), deadline(d) {

	visited.at(loc) = true;
}

std::vector<boost::optional<Op> > Rescue::ops() const {
	std::vector<boost::optional<Op> > ops;
	double left = deadline - time + Eps;
	if (left < ReportTime)
		goto skipreports;

	for (unsigned int i = 0; i < people.size(); i++) {
		if (people[i] == loc && !reported[i]) {
			ops.push_back(boost::optional<Op>(Op(Op::Report)));
			break;
		}
	}

skipreports:

	std::vector<int> ns = m.neighbors(loc);
	for (auto n = ns.begin(); n != ns.end(); n++) {
		const Map::Edge *e = m.edge(loc, *n);
		if (left >= e->line.length()) {
			assert(*n >= 0);
			ops.push_back(boost::optional<Op>(Op(Op::Move, *n)));
		}
	}

	if (loc == home)
		ops.push_back(boost::optional<Op>());

	return ops;
}

bool Rescue::consistent(const Rescue &o) const {
	return m == o.m && people == o.people && reported == o.reported;
}

void Rescue::act(boost::optional<Op> o) {
	if (!o) {
		time = deadline;
		return;
	}
	switch (o->kind) {
	case Op::Report: {
		bool r = false;
		for (unsigned int i = 0; i < people.size(); i++) {
			if (people[i] != loc)
				continue;
			r = true;
			reported[i] = true;
		}
		if (!r)
			fatal("Cannot report at %d: no people", loc);
		time += ReportTime;
		break;
	}

	case Op::Move:	{
		const Map::Edge *e = m.edge(loc, o->value);
		if (!e)
			fatal("Cannot move to %d from %d: no edge", o->value, loc);
		loc = o->value;
		visited.at(loc) = true;
		time += e->line.length();
		break;
	}

	default:
		fatal("Cannot perform operator type %d", o->kind);
	}

	if (time > deadline+Eps)
		fatal("Over the deadline!: time=%g, deadline=%g", time, deadline);
}

void Rescue::dump(FILE *out) const {
	fprintf(out, "nodes:\n");
	for (auto n = m.nodes.begin(); n != m.nodes.end(); n++) {
		if (n->id < 0)
			continue;
		fprintf(out, "\t%d:", n->id);
		for (int i = 0; i < Map::Ndirs; i++)
			fprintf(out, " %d", n->e[i]);
		fprintf(out, "\n");
	}
	fprintf(out, "edges:\n");
	for (auto e = m.edges.begin(); e != m.edges.end(); e++) {
		if (e->id < 0)
			continue;
		fprintf(out, "\t%d: %d↔%d, %g\n", e->id, e->n[0], e->n[1], e->line.length());
	}
	fprintf(out, "people:\n\t");
	for (unsigned int i = 0; i < people.size(); i++) {
		if (people[i] < 0)
			continue;
		if (reported[i])
			fprintf(out, " (%d)", people[i]);
		else
			fprintf(out, " %d", people[i]);
	}
	fprintf(out, "\n");
	fprintf(out, "loc=%d, home=%d\ntime=%g, deadline=%g\n",
		loc, home, time, deadline);
}

void Rescue::sense(const Rescue &w) {
	assert (w.loc == loc);

	assert (m.hasnode(loc));
	for (int i = 0; i < Map::Ndirs; i++) {
		int e = w.m.nodes[loc].e[i];
		m.nodes[loc].e[i] = e;
		if (e == Map::Node::MayEdge)
			m.nodes[loc].e[i] = Map::Node::NoEdge;
		if (e < 0)
			continue;

		int dst = w.m.edges[e].n[0];
		if (dst == loc)
			dst = w.m.edges[e].n[1];

		if (!m.hasnode(dst)) {
			Map::Node &nd = m.addnode(dst);
			nd.entrydir = w.m.nodes[dst].entrydir;
			nd.pt = w.m.nodes[dst].pt;
			nd.e[Map::rev(i)] = w.m.nodes[dst].e[Map::rev(i)];
		}

		if (!m.hasedge(e))
			m.addedge(e, loc, dst);
	}

	for (unsigned int i = 0; i < w.people.size(); i++) {
		if (w.people[i] != loc)
			continue;
		people[i] = w.people[i];
	}

	if (visited.size() < m.nodes.size())
		visited.resize(m.nodes.size(), false);

	if (w.ops() != ops()) {
		fprintf(stderr, "seen:\n");
		dump(stderr);
		fprintf(stderr, "ops:\n");
		std::vector<boost::optional<Op> > os = ops();
		for (auto o = os.begin(); o != os.end(); o++) {
			if (*o)
				fprintf(stderr, "\t%s\n", (*o)->str().c_str());
			else
				fprintf(stderr, "\tFinish\n");
		}

		fprintf(stderr, "\n\n----------\n");
		fprintf(stderr, "w:\n");
		w.dump(stderr);
		fprintf(stderr, "w.ops:\n");
		os = w.ops();
		for (auto o = os.begin(); o != os.end(); o++) {
			if (*o)
				fprintf(stderr, "\t%s\n", (*o)->str().c_str());
			else
				fprintf(stderr, "\tFinish\n");
		}
		assert (false);
	}
}

void Rescue::validateop(boost::optional<Op> op, std::vector<Op> oz) const {
	if (!op && loc != home) {
		draw("failed.eps", oz);
		fatal("Finish performed at the non end node");
	}
	if (!op)
		return;

	std::vector<boost::optional<Op> > os = ops();
	auto o = os.begin();
	for ( ; o != os.end(); o++) {
		if (**o == *op)
			break;
	}
	if (o == os.end()) {
		fprintf(stderr, "\n\n----------\nseen:\n");
		dump(stderr);
		fprintf(stderr, "ops:\n");
		std::vector<boost::optional<Op> > os = ops();
		for (auto o = os.begin(); o != os.end(); o++)
			fprintf(stderr, "\t%s\n", (*o)->str().c_str());
		fatal("Planner returned invalid operator %s\n", op->str().c_str());
	}
}

void Rescue::draw(std::string path, const std::vector<Op> &ops) const {
	Map map(m);
	for (auto p = people.begin(); p != people.end(); p++) {
		if (*p < 0)
			continue;
		map.nodes[*p].c = Image::red;
	}
	Map::Node *n = &map.nodes[home];
	n->c = Color(1, 0, 1);
	n->r = 5;
	for (auto o = ops.begin(); o != ops.end(); o++) {
		switch (o->kind) {
		case Op::Move: {
			int dst = o->value;
			int edge = -1;
			for (int j = 0; j < Map::Ndirs; j++) {
				if (n->e[j] < 0)
					continue;
				const Map::Edge &e = map.edges[n->e[j]];
				if (e.n[0] == dst || e.n[1] == dst) {
					edge = n->e[j];
					break;
				}
			}
			if (edge == -1)
				fatal("No edge from %d to %d", n->id, dst);
			map.edges[edge].c = Image::blue;
			map.edges[edge].w = 4;
			n = &map.nodes[dst];
			if (n->id != home)
				n->r = 3;
			break;
		}

		case Op::Report:
			n->c = Image::green;
			break;

		case Op::Skip:
			break;

		case Op::None:
			break;

		default:
			fatal("Unknown operator kind: %d", o->kind);
		}
	}

	map.nodes[loc].c = Color(1, 1, 0);
	map.nodes[loc].r = 5;

	map.draw(path.c_str());
}

unsigned int Rescue::npeople() const {
	unsigned int npeople = 0;
	for (auto p = people.begin(); p != people.end(); p++) {
		if (*p >= 0)
			npeople++;
	}
	return npeople;
}

Rescue Rescue::future(unsigned int sz) const {
	double PProb = 0.1;
	if (southonly + westonly > 1)
		PProb = 0.4;
	else if (southonly + westonly == 1)
		PProb = 0.2;

	Rescue f(*this);
	unsigned int n = f.m.nodes.size();
	f.m.grow(sz - f.m.nnodes());
	if (f.m.nodes.size() > f.visited.size())
		f.visited.resize(f.m.nodes.size(), false);
	unsigned int m = f.m.nodes.size();


	std::vector<int> newnodes(m - n);
	int j = 0;
	for (unsigned int i = n; i < m; i++)
		newnodes.at(j++) = i;
	shuf(newnodes);


	unsigned int added = npeople();
	for (unsigned int j = 0; j < newnodes.size(); j++) {
		int nd = newnodes.at(j);
		if (southonly && f.m.nodes[nd].pt.y >= 0)
			continue;
		if (westonly && f.m.nodes[nd].pt.x >= 0)
			continue;
		if (randgen.real() > PProb)
			continue;

		assert (f.m.nodes[nd].id >= 0);
		f.people.push_back(nd);
		f.reported.push_back(false);

		if (added+npeople() >= MaxPeople-1)
			break;
		added++;
	}
	return f;
}

Rescue Rescue::initallyknown() const {
	Map map;
	if (!map.hasnode(home))
		map.addnode(home);
	return Rescue(map, std::vector<int>(people.size(), -1), home, deadline);

}

void Rescue::done(FILE *out, const std::vector<Op>&) const {
	double cost = 0;
	unsigned int unreported = 0;
	for (auto r = reported.begin(); r != reported.end(); r++) {
		if (!*r) {
			unreported++;
			cost += RescueReward;
		}
	}

	unsigned int nvisited = 0;
	for (auto v = visited.begin(); v != visited.end(); v++) {
		if (*v)
			nvisited++;
	}

	dfpair(out, "map size (nodes)", "%u", (unsigned int) m.nodes.size());
	dfpair(out, "map size (edges)", "%u", (unsigned int) m.edges.size());
	dfpair(out, "nodes visited", "%u", nvisited);
	dfpair(out, "total people", "%u", (unsigned int) people.size());
	dfpair(out, "unreported people", "%u", unreported);
	dfpair(out, "deadline time", "%g", deadline);
	dfpair(out, "final execution time", "%g", time);
	dfpair(out, "final sol cost", "%g", cost);
}

class Search {
public:
	Search(const Graph&, int, int, double);
	Rescue::Solution search();
	void output(FILE*) const;

private:
	void dfs(int, double,std::vector<Graph::Shortest>&);

	// Problem instance.
	Graph g;
	int start, end;
	double deadline;
	std::vector<int> ign;	// Impossible people (ignored during search).

	// mintime is the minimum time for
	// reporting a group of people.  It is
	// indexed by the people ID, and it
	// includes the time to report and
	// then go, via the shortest path, to
	// the end.
	std::vector<double> mintime;

	// Statistics.
	double searchtime;
	unsigned long expd;	// Nodes expanded.

	// Solution.
	std::vector<int> solution;
	std::vector<int> skips;
	double cost, time;
	bool lastreport;	// True if the last move was to report;

	// Depth-first path.
	std::vector<int> path;	// Nodes visited.
	std::vector<int> left;	// Nodes with people left to report.
	unsigned int nl, depth;
};

Search::Search(const Graph &gr, int s, int e, double d) :
	g(gr), start(s), end(e), deadline(d),
	cost(std::numeric_limits<double>::infinity()) {

	g.precompute(start);
	g.precompute(end);
	for (auto p = g.people.begin(); p != g.people.end(); p++) {
		g.precompute(p->node);
		mintime.push_back(ReportTime+g.shortest[end][p->node].dist);
	}
}

Rescue::Solution Search::search() {
	double starttime = walltime();
	expd = 0;

	left.clear();
	ign.clear();
	for (unsigned int i = 0; i < g.people.size(); i++) {
		int n = g.people[i].node;
		if (g.shortest[start][n].dist + g.shortest[n][end].dist + ReportTime > deadline)
			ign.push_back(i);
		else
			left.push_back(n);
	}
	nl = left.size();
	path.clear();
	path.resize(g.people.size());
	depth = 0;
	dfs(start, 0, g.shortest[start]);

	std::vector<Op> ops;
	int cur = start;
	for (auto n = solution.begin(); n != solution.end(); n++) {
		Graph::Shortest *s = &g.shortest[*n][cur];
		while (s->parent >= 0) {
			ops.push_back(Op(Op::Move, s->parent));
			s = &g.shortest[*n][s->parent];
		}
		cur = *n;
		ops.push_back(Op(Op::Report));
	}
	if (ops.size() > 0 && !lastreport)
		ops.pop_back();
	for (auto s = skips.begin(); s != skips.end(); s++)
		ops.push_back(Op(Op::Skip, *s));
	for (auto s = ign.begin(); s != ign.end(); s++) {
		ops.push_back(Op(Op::Skip, *s));
		cost += g.people[*s].num;
	}

	searchtime = walltime() - starttime;
	Rescue::Solution s;
	s.ops = ops;
	s.ct = CostTime(cost, time);
	return s;
}

void Search::dfs(int loc, double t, std::vector<Graph::Shortest> &s) {
	assert (t <= deadline);

	double h = 0;
	for (unsigned int i = 0; i < nl; i++) {
		int dst = left[i];
		if (t+s[dst].dist+mintime[g.nodes[dst].p] <= deadline)
			continue;
		h += RescueReward*g.people[g.nodes[dst].p].num;
		if (h >= cost)
			return;
	}

	expd++;

	// Consider reporting each remaining room.
	for (unsigned int i = 0; i < nl; i++) {
		int dst = left[i];
		if (t+s[dst].dist+mintime[g.nodes[dst].p] > deadline)
			continue;
		left[i] = left[nl-1];
		nl--;
		path[depth++] = dst;

		dfs(dst, t+s[dst].dist+ReportTime, g.shortest[dst]);

		depth--;
		nl++;
		left[nl-1] = left[i];
		left[i] = dst;
	}

	// Consider just ending now.
	if (s[end].dist > deadline)
		return;

	double c = 0;
	for (unsigned int i = 0; i < nl; i++) {
		c += RescueReward*g.people[g.nodes[left[i]].p].num;
		if (c >= cost)
			return;
	}

	solution.clear();
	std::copy_n(path.begin(), depth, std::back_inserter(solution));

	lastreport = false;
	if (loc != end) {
		solution.push_back(end);
	} else if (depth > 0) {
		// Must have moved here first to report someone
		// then decided to end it.
		lastreport = true;
	}

	skips.clear();
	std::copy_n(left.begin(), nl, std::back_inserter(skips));
	time = t + s[end].dist;
	cost = c;
}

void Search::output(FILE *o) const {
	dfpair(o, "deadline", "%g", deadline);
	dfpair(o, "execution time", "%g", time);
	dfpair(o, "final sol cost", "%g", cost);
	dfpair(o, "final sol length", "%u", (unsigned int) solution.size());
	dfpair(o, "ignored occupied rooms", "%u", (unsigned int) ign.size());
	dfpair(o, "number of skipped rooms", "%u", (unsigned int) skips.size());
	dfpair(o, "total wall time", "%g", searchtime);
	dfpair(o, "total nodes expanded", "%lu", expd);
}

Rescue::Solution Rescue::solve() const {
	Graph g = m.rescuegraph();
	for (unsigned int i = 0; i < people.size(); i++) {
		if (people[i] >= 0 && !reported[i])
			g.addperson(people[i]);
	}
	Search dfs(g, loc, home, deadline-time);
	auto sol = dfs.search();
	if (sol.ops.empty())
		return sol;
	while (!sol.ops.empty() && sol.ops.back().kind != Op::Move && sol.ops.back().kind != Op::Report)
		sol.ops.pop_back();
	return sol;
}
